from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException
from jose import jwt, JWTError
from starlette import status
from starlette.middleware.cors import CORSMiddleware
from starlette.middleware.exceptions import ExceptionMiddleware
from starlette.requests import Request
from starlette.responses import PlainTextResponse, JSONResponse

from api import sys, case, R


async def startup():
    pass


async def shutdown():
    pass


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load the ML model
    await startup()
    yield
    # Clean up the ML models and release the resources
    await shutdown()


app = FastAPI(lifespan=lifespan)

secret_key = b'6dfaa67b9e9211eca1b40242ac110002'


def decode_access_token(token: str):
    try:
        payload = jwt.decode(token, secret_key, algorithms=['HS256'])
        return payload
    except JWTError as e:
        print(e)
        return None


@app.middleware("http")
async def auth_middleware(request: Request, call_next):
    auth_header = request.headers.get("Auth")
    if auth_header:
        # extract token from the header and keep it in the global variable
        auth_token = auth_header.split(" ")[0]
        payload = decode_access_token(auth_token)
        if not payload:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Invalid authentication credentials",
                headers={"WWW-Authenticate": "Bearer"},
            )
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="authentication credentials is empty",
            headers={"WWW-Authenticate": "Bearer"},
        )
    response = await call_next(request)
    return response

@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    print(exc)
    return JSONResponse(status_code=exc.status_code, content=R.fail(msg=exc.detail))


# 添加 CORS 中间件（可选）
app.add_middleware(
    CORSMiddleware,
    allow_origins=['http://localhost:9000', 'https://ys-tiaojie.hnlabour.com'],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.add_middleware(ExceptionMiddleware, handlers=app.exception_handlers)

app.include_router(sys.router)
app.include_router(case.router)
if __name__ == '__main__':
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=10001)
